import numpy as np 
import h5py

class ExperienceCollector:
    def __init__(self) -> None:
        self.states = []
        self.actions = []
        self.rewards = []
        self.advantages = []
        self.current_episode_states = []
        self.current_episode_actions = []
        self.current_episode_estimated_values = []

    def BeginEpisode(self):
        self.current_episode_actions = []
        self.current_episode_estimated_values = []
        self.current_episode_states = []

    def RecordDecision(self, state, action, estimated_value=0):
        self.current_episode_states.append(state)
        self.current_episode_actions.append(action)
        self.current_episode_estimated_values.append(estimated_value)

    def CompleteEpisode(self, reward):
        num_states = len(self.current_episode_states)
        self.states += self.current_episode_states
        self.actions += self.current_episode_actions
        self.rewards += [reward for _ in range(num_states)]
        for i in range(num_states):
            advantage = reward - self.current_episode_estimated_values[i]
            self.advantages.append(advantage)
        self.current_episode_actions = []
        self.current_episode_states = []
        self.current_episode_estimated_values = []

class ExperienceBuffer:
    def __init__(self, states, actions, rewards, advantages) -> None:
        self.states = states
        self.actions = actions
        self.rewards = rewards
        self.advantages = advantages

    def Serialize(self, f: h5py.File):
        ex = 'experience'
        f.create_group(ex)
        f[ex].create_dataset('states', data=self.states)
        f[ex].create_dataset('actions', data=self.actions)
        f[ex].create_dataset('rewards', data=self.rewards)
        f[ex].create_dataset('advantages', data=self.advantages)

def CombineExperience(collectors):
    states = np.concatenate([np.array(c.states) for c in collectors])
    actions = np.concatenate([np.array(c.actions) for c in collectors])
    rewards = np.concatenate([np.array(c.rewards) for c in collectors])
    advantages = np.concatenate([np.array(c.advantages) for c in collectors])
    return ExperienceBuffer(states, actions, rewards, advantages)

def LoadExperience(f: h5py.File):
    ex = 'experience'
    states = np.array(f[ex]['states'])
    actions = np.array(f[ex]['actions'])
    rewards = np.array(f[ex]['rewards'])
    advantages = np.array(f[ex]['advantages'])
    return ExperienceBuffer(states, actions, rewards, advantages)